#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Rating prediction functions for collaborative filtering recommendation system.

Created on: 2025-04-18
Author: Nianqing Liu
"""

import numpy as np


def simple_weighted_average(ratings_matrix, similarity_matrix, k=10):
    """
    Predict ratings using simple weighted average of similar users' ratings.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        User-item rating matrix.
    similarity_matrix : numpy.ndarray
        User-user similarity matrix.
    k : int
        Number of similar users to consider (default 10).

    Returns:
    --------
    numpy.ndarray
        Predicted rating matrix.
    """
    n_users, n_items = ratings_matrix.shape
    predicted_ratings = np.zeros((n_users, n_items))

    # For each user
    for u in range(n_users):
        # Find k most similar users (excluding self)
        user_similarities = similarity_matrix[u]
        # Set similarity with self to -1 to exclude it
        user_similarities[u] = -1
        similar_users = np.argsort(user_similarities)[::-1][:k]

        # For each item
        for i in range(n_items):
            # Get similar users who have rated this item
            sim_users_rated = [v for v in similar_users if ratings_matrix[v, i] > 0]

            # If no similar user has rated the item, skip
            if len(sim_users_rated) == 0:
                predicted_ratings[u, i] = 0
                continue

            # Calculate weighted average
            sim_sum = sum(abs(similarity_matrix[u, v]) for v in sim_users_rated)
            if sim_sum == 0:
                predicted_ratings[u, i] = 0
                continue

            weighted_sum = sum(
                similarity_matrix[u, v] * ratings_matrix[v, i] for v in sim_users_rated
            )

            predicted_ratings[u, i] = weighted_sum / sim_sum

    return predicted_ratings


def bias_weighted_average(ratings_matrix, similarity_matrix, k=10):
    """
    Predict ratings using weighted average with user bias correction.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        User-item rating matrix.
    similarity_matrix : numpy.ndarray
        User-user similarity matrix.
    k : int
        Number of similar users to consider (default 10).

    Returns:
    --------
    numpy.ndarray
        Predicted rating matrix.
    """
    n_users, n_items = ratings_matrix.shape
    predicted_ratings = np.zeros((n_users, n_items))

    # Calculate user mean ratings
    user_rated_mask = ratings_matrix > 0
    user_ratings_count = np.sum(user_rated_mask, axis=1)
    user_ratings_sum = np.sum(ratings_matrix, axis=1)

    # Handle division by zero with np.where
    user_mean_ratings = np.where(
        user_ratings_count > 0, user_ratings_sum / user_ratings_count, 0
    )

    # For each user
    for u in range(n_users):
        # Find k most similar users (excluding self)
        user_similarities = similarity_matrix[u]
        # Set similarity with self to -1 to exclude it
        user_similarities[u] = -1
        similar_users = np.argsort(user_similarities)[::-1][:k]

        # User's mean rating
        u_mean = user_mean_ratings[u]

        # For each item
        for i in range(n_items):
            # Skip if user has already rated this item
            if ratings_matrix[u, i] > 0:
                predicted_ratings[u, i] = ratings_matrix[u, i]
                continue

            # Get similar users who have rated this item
            sim_users_rated = [v for v in similar_users if ratings_matrix[v, i] > 0]

            # If no similar user has rated the item, use user's mean rating or global mean
            if len(sim_users_rated) == 0:
                predicted_ratings[u, i] = (
                    u_mean
                    if u_mean > 0
                    else np.mean(ratings_matrix[ratings_matrix > 0])
                )
                continue

            # Calculate bias-adjusted weighted average
            sim_sum = sum(abs(similarity_matrix[u, v]) for v in sim_users_rated)
            if sim_sum == 0:
                predicted_ratings[u, i] = u_mean
                continue

            weighted_sum = sum(
                similarity_matrix[u, v] * (ratings_matrix[v, i] - user_mean_ratings[v])
                for v in sim_users_rated
            )

            predicted_rating = u_mean + weighted_sum / sim_sum

            # Clip the predicted rating to the valid range [1,5]
            predicted_ratings[u, i] = max(1, min(5, predicted_rating))

    return predicted_ratings


def predict_ratings(ratings_matrix, similarity_matrix, method="bias_weighted", k=10):
    """
    Predict ratings using the specified method.

    Parameters:
    -----------
    ratings_matrix : numpy.ndarray
        User-item rating matrix.
    similarity_matrix : numpy.ndarray
        User-user similarity matrix.
    method : str
        Prediction method: 'simple_weighted' or 'bias_weighted'.
    k : int
        Number of similar users to consider.

    Returns:
    --------
    numpy.ndarray
        Predicted rating matrix.
    """
    if method == "simple_weighted":
        return simple_weighted_average(ratings_matrix, similarity_matrix, k)
    elif method == "bias_weighted":
        return bias_weighted_average(ratings_matrix, similarity_matrix, k)
    else:
        raise ValueError(
            f"Unknown prediction method: {method}. Choose from 'simple_weighted' or 'bias_weighted'."
        )
